/*
 * Copyright (c) 2018-2023 北京华恒盛世科技有限公司
 * QUICKPOOL is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *     http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
*/

#include "subcmd.h"

#define SIGCHK   -1
#define SIGFORCE -2

extern char *sigsymbol_2string(void);

static void print_message(int, JOBID_T);
static void signaljob_usage(const char *, int);

/**
 * @brief       qp job kill/stop/resume/chkpnt handler, to send signal to a job.
 *
 * @param[in]   argc          #1: number of input argument
 * @param[in]   argv          #2: string array of the input argument
 *
 * @retval      0             #1: succeed
 * @retval      1             #2: input option is wrong or kind of input argument not specified
 * @retval      2             #3: argument has syntax error
 * @retval      3             #4: failed to send signal to the job
 */
int
cmd_signal_job(int argc, char **argv)
{
    int signalValue=0;

    const char *options;
    static struct option long_options[10];
    long_options[0].name = "help";
    long_options[0].has_arg = no_argument;
    long_options[0].flag = 0;
    long_options[0].val = 'h';

    long_options[1].name = "node";
    long_options[1].has_arg = required_argument;
    long_options[1].flag = 0;
    long_options[1].val = 'm';

    long_options[2].name = "queue";
    long_options[2].has_arg = required_argument;
    long_options[2].flag = 0;
    long_options[2].val = 'q';

    long_options[3].name = "user";
    long_options[3].has_arg = required_argument;
    long_options[3].flag = 0;
    long_options[3].val = 'u';

    long_options[4].name = "jobname";
    long_options[4].has_arg = required_argument;
    long_options[4].flag = 0;
    long_options[4].val = 'J';

    if (strstr(argv[0], "kill")) {

        signalValue = SIGKILL;
        options = "hlm:q:u:J:g:rs:";

        long_options[5].name = "group";
        long_options[5].has_arg = required_argument;
        long_options[5].flag = 0;
        long_options[5].val = 'g';

        long_options[6].name = "list";
        long_options[6].has_arg = no_argument;
        long_options[6].flag = 0;
        long_options[6].val = 'l';

        long_options[7].name = "force";
        long_options[7].has_arg = no_argument;
        long_options[7].flag = 0;
        long_options[7].val = 'r';

        long_options[8].name = "signal";
        long_options[8].has_arg = required_argument;
        long_options[8].flag = 0;
        long_options[8].val = 's';

        long_options[9].name = NULL;
        long_options[9].has_arg = no_argument;
        long_options[9].flag = 0;
        long_options[9].val = 0;
    } else if (strstr(argv[0], "stop")) {

        signalValue = SIGSTOP;
        options = "hm:q:u:J:g:";

        long_options[5].name = "group";
        long_options[5].has_arg = required_argument;
        long_options[5].flag = 0;
        long_options[5].val = 'g';

        long_options[6].name = NULL;
        long_options[6].has_arg = no_argument;
        long_options[6].flag = 0;
        long_options[6].val = 0;
    } else if (strstr(argv[0], "resume")) {

        signalValue = SIGCONT;
        options = "hm:q:u:J:g:";

        long_options[5].name = "group";
        long_options[5].has_arg = required_argument;
        long_options[5].flag = 0;
        long_options[5].val = 'g';

        long_options[6].name = NULL;
        long_options[6].has_arg = no_argument;
        long_options[6].flag = 0;
        long_options[6].val = 0;
    } else if (strstr(argv[0], "chkpnt")) {

        signalValue = SIGCHK;    
        options = "hm:q:u:J:p:fkS";

        long_options[5].name = "period";
        long_options[5].has_arg = required_argument;
        long_options[5].flag = 0;
        long_options[5].val = 'p';

        long_options[6].name = "force";
        long_options[6].has_arg = no_argument;
        long_options[6].flag = 0;
        long_options[6].val = 'f';

        long_options[7].name = "kill";
        long_options[7].has_arg = no_argument;
        long_options[7].flag = 0;
        long_options[7].val = 'k';

        long_options[8].name = "stop";
        long_options[8].has_arg = no_argument;
        long_options[8].flag = 0;
        long_options[8].val = 'S';

        long_options[9].name = NULL;
        long_options[9].has_arg = no_argument;
        long_options[9].flag = 0;
        long_options[9].val = 0;
    } else {
        fprintf(stderr, "Unknown operation %s.\n", argv[0]);
        exit(2);
    }

    char *strUser=NULL, *strQueue=NULL, *strNode=NULL, *strJobName=NULL, *strJobGroup=NULL;
    int chkpntPeriod=0, chkpntOptions=0;
    int sigValue = signalValue;

    bool_t kOpt=FALSE, SOpt=FALSE, rOpt=FALSE, sOpt=FALSE;
    while (TRUE) {
        int option_index = 0;
        int opt;
        opt = getopt_long(argc, argv, options, long_options, &option_index);
        if (-1 == opt) {
            break;
        }

        switch (opt) {
        case 'm':
            if (NULL != strNode) {
                fprintf(stderr, "Do not specify -m/--node option for multiple times.\n");
                signaljob_usage(argv[0], signalValue);
                exit(1);
            }

            if (0 >= strlen(optarg)) {
                fprintf(stderr, "Specify a node name with -m/--node option.\n");
                signaljob_usage(argv[0], signalValue);
                exit(1);
            }

            strNode = optarg;
            break;
        case 'q':
            if (NULL != strQueue) {
                fprintf(stderr, "Do not specify -q/--queue option for multiple times.\n");
                signaljob_usage(argv[0], signalValue);
                exit(1);
            }

            if (0 >= strlen(optarg)) {
                fprintf(stderr, "Specify a queue name with -q/--queue option.\n");
                signaljob_usage(argv[0], signalValue);
                exit(1);
            }

            strQueue = optarg;
            break;
        case 'u':
            if (NULL != strUser) {
                fprintf(stderr, "Do not specify -u/--user option for multiple times.\n");
                signaljob_usage(argv[0], signalValue);
                exit(1);
            }

            if (0 >= strlen(optarg)) {
                fprintf(stderr, "Specify a user name with -u/--user option.\n");
                signaljob_usage(argv[0], signalValue);
                exit(1);
            }

            strUser = optarg;
            break;
        case 'J':
            if (NULL != strJobName) {
                fprintf(stderr, "Do not specify -J/--jobname option for multiple times.\n");
                signaljob_usage(argv[0], signalValue);
                exit(1);
            }

            if (0 >= strlen(optarg)) {
                fprintf(stderr, "Specify a job name with -J/--jobname option.\n");
                signaljob_usage(argv[0], signalValue);
                exit(1);
            }

            strJobName = optarg;
            break;
        case 'g':
            if (NULL != strJobGroup) {
                fprintf(stderr, "Do not specify -g/--jobgroup option for multiple times.\n");
                signaljob_usage(argv[0], signalValue);
                exit(1);
            }

            if (0 >= strlen(optarg)) {
                fprintf(stderr, "Specify a job group name with -g/--jobgroup option.\n");
                signaljob_usage(argv[0], signalValue);
                exit(1);
            }

            strJobGroup = optarg;
            break;
        case 'l':
            fprintf(stdout, "%s", sigsymbol_2string());
            exit (0);
        case 'f':
            chkpntOptions |= QP_CHKPNT_FORCE;
            break; 
        case 's':
            if (rOpt) {
                fprintf(stderr, "%s: -s/--signal option does not work with -r/--force.\n", argv[0]);
                signaljob_usage(argv[0], signalValue);
                exit(1);
            }

            if ((sigValue = gf_string_2signal(optarg)) < 0) {
                fprintf(stderr, "Specified unknown signal value %s for -s/--signal option.\n", optarg);
                exit(2);
            }

            sOpt = TRUE;

            break;
        case 'r':
            if (sOpt) {
                fprintf(stderr, "%s: -r/--force option does not work with -s/--signal.\n", argv[0]);
                signaljob_usage(argv[0], signalValue);
                exit(1);
            }

            sigValue = SIGFORCE;
            rOpt = TRUE;

            break;
        case 'p':
            chkpntPeriod = gf_atoi(optarg, INFINIT_INT, 0);
            if (INFINIT_INT == chkpntPeriod) {
                fprintf(stderr, "%s: Must specify a positive integer for -p/--period option.", optarg);
                exit(2);
            }

            break;
        case 'k':
            if (SOpt) {
                fprintf(stderr, "%s: -k/--kill option does not work with -S/--stop.\n", argv[0]);
                signaljob_usage(argv[0], signalValue);
                exit(1);
            }

            kOpt = TRUE;
            chkpntOptions |= QP_CHKPNT_KILL;

            break;
        case 'S':
            if (kOpt) {
                fprintf(stderr, "%s: -S/--stop option does not work with -k/--kill.\n", argv[0]);
                signaljob_usage(argv[0], signalValue);
                exit(1);
            }

            SOpt = TRUE;
            chkpntOptions |= QP_CHKPNT_STOP;

            break;
        case 'h':
            signaljob_usage(argv[0], signalValue);
            exit(0);
        case '?':
        default:
            signaljob_usage(argv[0], signalValue);
            exit(1);
        }
    }

    if (argc == optind && NULL == strUser && NULL == strQueue
        && NULL == strNode && NULL == strJobName && NULL == strJobGroup) {

        fprintf(stderr, "Job ID or one of '-m/--node', '-u/--user' '-q/--queue' '-g/--jobgroup' and '-J/--jobname' must be specified.\n");
        signaljob_usage(argv[0], signalValue);

        exit(2);
    }

    if ((SIGKILL == signalValue) && (SIGFORCE != sigValue)
        && (sigValue >= NSIG || sigValue < 1)) {

        fprintf(stderr, "Specified unknown signal value <%d>.\n", sigValue);
        exit(2);
    }


    struct request signalRequest;
    memset(&signalRequest, 0, sizeof(struct request));

    unsigned int num = 0;

    int numJobIds;
    JOBID_T *jobIds = NULL;
    numJobIds = get_jobids_fromargv(argc, argv, &jobIds);
    if (0 < numJobIds) { // specified job id, ignore -m, -u, -q, -g and -J options.

        num += 2; //jobIDs, sigValue

        if (SIGCHK == sigValue) {
            if (0 < chkpntOptions) {
                num ++;
            }
            if (0 < chkpntPeriod) {
                num ++;
            }
        }

        signalRequest.number = num;
        signalRequest.keyvalues = (struct keyvalue *)calloc(num, sizeof(struct keyvalue));
        if (NULL == signalRequest.keyvalues) {
            fprintf(stderr, "calloc(%ld) failed due to %m.\n",
                    signalRequest.number*sizeof(struct keyvalue));
            exit(4);
        }

        num = 0;

        signalRequest.keyvalues[num].key = REQUEST_KEY_SIGNAL;
        signalRequest.keyvalues[num].value = (char *)calloc(12, sizeof(char));
        if (NULL == signalRequest.keyvalues[num].value) {
            fprintf(stderr, "calloc(%d) failed due to %m.\n", 12);
            exit(4);
        }
        if (SIGCHK == sigValue) {
            snprintf(signalRequest.keyvalues[num].value, 12, "%d", QUICKPOOL_SIGCHKPNT);
        } else if (SIGFORCE == sigValue) {
            snprintf(signalRequest.keyvalues[num].value, 12, "%d", QUICKPOOL_SIGKILL_FORCE);
        } else {
            snprintf(signalRequest.keyvalues[num].value, 12, "%d", sigValue);
        }
        num ++;

        unsigned int size = numJobIds*sizeof(JOBID_T) + 10;
        signalRequest.keyvalues[num].key = REQUEST_KEY_JOBIDS;
        signalRequest.keyvalues[num].value = (char *)calloc(size+1, sizeof(char));
        if (NULL == signalRequest.keyvalues[num].value) {
            fprintf(stderr, "calloc(%d) failed due to %m.\n", size+1);
            exit(4);
        }

        int i;
        for (i = 0; i < numJobIds; i++) {
            if (0 == i) {
                snprintf(signalRequest.keyvalues[num].value, size, "%lld", jobIds[i]);
                continue;
            }

            unsigned int lenValue;
            lenValue = strlen(signalRequest.keyvalues[num].value);
            snprintf(signalRequest.keyvalues[num].value+lenValue, size, " %lld", jobIds[i]);
        }
        num ++;

        FREEUP(jobIds);

        if (SIGCHK == sigValue) {
            if (0 < chkpntOptions) {
                signalRequest.keyvalues[num].key = REQUEST_KEY_CHKPNT_FLAG;
                signalRequest.keyvalues[num].value = (char *)calloc(12, sizeof(char));
                if (NULL == signalRequest.keyvalues[num].value) {
                    fprintf(stderr, "calloc(%d) failed due to %m.\n", 12);
                    exit(4);
                }
                snprintf(signalRequest.keyvalues[num].value, 12, "%d", chkpntOptions);

                num ++;
            }
            if (0 < chkpntPeriod) {
                signalRequest.keyvalues[num].key = REQUEST_KEY_CHKPNT_PERIOD;
                signalRequest.keyvalues[num].value = (char *)calloc(12, sizeof(char));
                if (NULL == signalRequest.keyvalues[num].value) {
                    fprintf(stderr, "calloc(%d) failed due to %m.\n", 12);
                    exit(4);
                }
                snprintf(signalRequest.keyvalues[num].value, 12, "%d", chkpntPeriod);

                num ++;
            }
        }
    } else {

        if (SIGKILL != sigValue && SIGSTOP != sigValue && SIGCONT != sigValue
            && SIGFORCE != sigValue) {

            fprintf(stderr, "You can only kill, stop or resume multiple jobs.\n");
            signaljob_usage(argv[0], signalValue);
            exit(2);
        }

        num ++; //sigValue

        if (NULL != strQueue) {
            num ++;
        }

        if (NULL != strNode) {
            num ++;
        }

        num ++; //username

        if (NULL != strJobName) {
            num ++;
        }

        if (NULL != strJobGroup) {
            num ++;
        }

        signalRequest.number = num;
        signalRequest.keyvalues = (struct keyvalue *)calloc(num, sizeof(struct keyvalue));
        if (NULL == signalRequest.keyvalues) {
            fprintf(stderr, "calloc(%ld) failed due to %m.\n",
                    signalRequest.number*sizeof(struct keyvalue));
            exit(4);
        }

        num = 0;

        signalRequest.keyvalues[num].key = REQUEST_KEY_SIGNAL;
        signalRequest.keyvalues[num].value = (char *)calloc(12, sizeof(char));
        if (NULL == signalRequest.keyvalues[num].value) {
            fprintf(stderr, "calloc(%d) failed due to %m.\n", 12);
            exit(4);
        }
        if (SIGCHK == sigValue) {
            snprintf(signalRequest.keyvalues[num].value, 12, "%d", QUICKPOOL_SIGCHKPNT);
        } else if (SIGFORCE == sigValue) {
            snprintf(signalRequest.keyvalues[num].value, 12, "%d", QUICKPOOL_SIGKILL_FORCE);
        } else {
            snprintf(signalRequest.keyvalues[num].value, 12, "%d", sigValue);
        }
        num ++;

        if (NULL != strQueue) {
            signalRequest.keyvalues[num].key = REQUEST_KEY_QUEUENAME;
            signalRequest.keyvalues[num].value = copy_string(strQueue);

            num ++;
        }

        if (NULL != strNode) {
            signalRequest.keyvalues[num].key = REQUEST_KEY_HOSTNAME;
            signalRequest.keyvalues[num].value = copy_string(strNode);

            num ++;
        }

        signalRequest.keyvalues[num].key = REQUEST_KEY_USERNAME;
        if (NULL == strUser) {
            signalRequest.keyvalues[num].value = getmy_uname();
        } else {
            signalRequest.keyvalues[num].value = copy_string(strUser);
        }
        num ++;

        if (NULL != strJobName) {
            signalRequest.keyvalues[num].key = REQUEST_KEY_JOBNAME;
            signalRequest.keyvalues[num].value = copy_string(strJobName);

            num ++;
        }

        if (NULL != strJobGroup) {
            signalRequest.keyvalues[num].key = REQUEST_KEY_SUBMIT_JOBGROUP;
            signalRequest.keyvalues[num].value = copy_string(strJobGroup);

            num ++;
        }
    }

    const char *op;
    switch (sigValue) {
    case SIGSTOP:
        op = "stop";
        break;
    case SIGCONT:
        op = "resume";
        break;
    case SIGKILL:
    case SIGFORCE:
        op = "terminate";
        break;
    default:
        op = "signal";
        break;
    }

    struct QuickPoolJobReply *jobReply;
    jobReply = client_signaljob(&signalRequest, NULL);
    if (NULL == jobReply) {

        gf_xdr_free(gf_xdr_request, &signalRequest);
        fprintf(stderr, "Faild to %s job due to above reason.\n", op);

        exit(3);
    }

    gf_xdr_free(gf_xdr_request, &signalRequest);

    unsigned int i;
    for (i = 0; i < jobReply->num_job; i++) {
        print_message(sigValue, jobReply->jobIDs[i]);
    }

    for (i = 0; i < jobReply->num_error; i++) {
        fprintf(stderr, "%s\n", jobReply->error_messages[i]);
    }

    if (0 < jobReply->num_error) {
        gf_xdr_free(libxdr_job_reply, jobReply);
        FREEUP(jobReply);
        exit(3);
    }

    gf_xdr_free(libxdr_job_reply, jobReply);
    FREEUP(jobReply);

    exit(0);
} // end function cmd_signal_job

/* ------------------------------------------------
 *
 *  below are static functions used in this file
 *
 * ------------------------------------------------ */

static void
print_message(int sig, JOBID_T jobid)
{
    const char *op;
    switch (sig) {
    case SIGSTOP:
        op = "stopped";
        break;
    case SIGCONT:
        op = "resumed";
        break;
    case SIGKILL:
    case SIGFORCE:
        op = "terminated";
        break;
    default:
        op = "signaled";
        break;
    }

    char strJobId[32];
    printf("Job <%s> is being %s\n", jobid_2string(jobid, strJobId), op);

    return;
} // end function print_message

static void
signaljob_usage(const char *cmd, int sig)
{
    fprintf(stdout, "qp job %s usage: ", cmd);
    fprintf(stdout, "[-h, --help] ");

    if (sig == SIGKILL) {
        fprintf(stdout, "[-l, --list]\n");
        fprintf(stdout, "\t\t   [-m, --node host_name] [-q, --queue queue_name] ");
        fprintf(stdout, "[-g, --group group_name]\n");
        fprintf(stdout, "\t\t   [-u, --user user_name|all] [-r, --force | -s, --signal sigval]\n");
        fprintf(stdout, "\t\t   [-J, --jobname job_name] [0 | jobId | \"jobId[index_list]\"]...\n");
    } else if (sig == SIGSTOP) {
        fprintf(stdout, "\n\t\t   [-m, --node host_name] [-q, --queue queue_name] ");
        fprintf(stdout, "[-g, --group group_name]\n");
        fprintf(stdout, "\t\t   [-u, --user user_name|all] [-J, --jobname job_name]\n");
        fprintf(stdout, "\t\t   [0 | jobId | \"jobId[index_list]\"]...\n");
    } else if (sig == SIGCONT) {
        fprintf(stdout, "\n\t\t     [-m, --node host_name] [-q, --queue queue_name] ");
        fprintf(stdout, "[-g, --group group_name]\n");
        fprintf(stdout, "\t\t     [-u, --user user_name|all] [-J, --jobname job_name]\n");
        fprintf(stdout, "\t\t     [0 | jobId | \"jobId[index_list]\"]...\n");
    } else if (sig == SIGCHK) {
        fprintf(stdout,  "\n\t\t     [-m, --node host_name] [-q, --queue queue_name] ");
        fprintf(stdout, "[-u, --user user_name|all]\n");
        fprintf(stdout, "\t\t     [-p, --period period] [-f, --force] [-k, --kill | -S, --stop]\n");
        fprintf(stdout, "\t\t     [-J, --jobname job_name] [jobId | \"jobId[index_list]\"]...\n");
    }

    return;
} // end function signaljob_usage
